from llama_index.core import VectorStoreIndex,SimpleDirectoryReader,Settings
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.huggingface import HuggingFaceLLM

def main():
    # 初始化一个HuggingFaceEmbedding对象，用于将文本转换为向量表示。
    # 指定一个预训练的sentence-transformers模型，这里使用"sentence-transformers/all-MiniLM-L6-v2"模型
    embed_model = HuggingFaceEmbedding(model_name=r"D:\self\python\AIModel\paraphrase-multilingual-MiniLM-L12-v2")

    # 将创建的嵌入模型赋值给全局设置的embed_model属性

    Settings.embed_model = embed_model

    llm = HuggingFaceLLM(model_name=r"D:\self\python\AIModel\Qwen2.5-0.5B-Instruct",tokenizer_name=r"D:\self\python\AIModel\Qwen2.5-0.5B-Instruct",model_kwargs={"trust_remote_code":True},tokenizer_kwargs={"trust_remote_code":True})

    #设置全局llm属性，这样在所以查询时，会使用该模型
    Settings.llm = llm

    # RAG 系统构建过程
    # 从指定目录读取所有文档，将加载到内存中，requeired_exts只加载指定扩展名的文档
    documents = SimpleDirectoryReader("./data",required_exts=[".md",".pdf"]).load_data()

    # 创建一个VectorStoreIndex对象，并将文档作为参数传入构建索引，次索引将文档转换为向量，默认放在内存中。
    index = VectorStoreIndex.from_documents(documents,show_progress=True)

    # 创建一个查询引擎，并将索引作为参数传入，查询引擎将使用索引来回答问题。
    query_engine = index.as_query_engine()
    response = query_engine.query("ai web框架")

    print(response)




if __name__ == '__main__':
    main()